import unittest
import pandas as pd
from analysis.analysis_falldown import find_all_fall

class TestFindAllFall(unittest.TestCase):
    def setUp(self):
        # 在每个测试方法之前设置测试数据
        self.test_data1 = {
            'date': ['2022-01-04', '2022-01-05', '2022-01-06', '2022-01-07'],
            'code': ['sh.600000'] * 4,
            'open': [100, 101, 100, 90],
            'close': [99, 100, 95, 85],
            'preclose': [101, 99, 100, 95]
        }
        self.test_data2 = {
            'date': ['2022-02-04', '2022-02-05', '2022-02-06', '2022-02-07'],
            'code': ['sh.600000'] * 4,
            'open': [200, 201, 200, 190],
            'close': [202, 201, 205, 210],
            'preclose': [200, 201, 200, 200]
        }
        self.dates = pd.to_datetime(['2022-01-04', '2022-01-05', '2022-01-06', '2022-01-07'])
        self.codes = pd.Series(['sh.600000'] * 4)

    def test_satisfying_conditions(self):
        # 创建DataFrame并计算跌幅
        df1 = pd.DataFrame(self.test_data1)
        # 调用函数并获取结果
        result = find_all_fall(df1)
        self.assertTrue(bool(result))
        self.assertEqual(result[0], result[1])

    def test_not_satisfying_conditions(self):
        # 创建DataFrame并计算跌幅
        df2 = pd.DataFrame(self.test_data2)
        # 调用函数并验证结果为空列表
        result = find_all_fall(df2)
        self.assertEqual(result[0], result[1])

    def tearDown(self):
        # 在每个测试方法之后清理测试数据
        del self.test_data1, self.test_data2, self.dates, self.codes

if __name__ == '__main__':
    unittest.main()